﻿#include "AIEnemyController.hpp"
#include <cmath>
#include <Windows.h>
#include <filesystem>

static std::wstring ToWideString(const std::string& str) {
    int size = MultiByteToWideChar(CP_UTF8, 0, str.c_str(), (int)str.size(), nullptr, 0);
    std::wstring wide(size, 0);
    MultiByteToWideChar(CP_UTF8, 0, str.c_str(), (int)str.size(), &wide[0], size);
    return wide;
}

AIEnemyController::AIEnemyController(const std::string& modelPath)
    : env(ORT_LOGGING_LEVEL_WARNING, "EnemyAI") {

    //std::cout << "Current working directory: " << std::filesystem::current_path() << std::endl;

    if (!std::filesystem::exists(modelPath)) {
        std::cerr << "Error: Model file not found " << modelPath << std::endl;
        throw std::runtime_error("Model file missing");
    }

    Ort::SessionOptions options;
    options.SetIntraOpNumThreads(1);
    options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);

    try {
        std::wstring widePath = ToWideString(modelPath);
        session = std::make_unique<Ort::Session>(env, widePath.c_str(), options);
    }
    catch (const Ort::Exception& e) {
        std::cerr << "ONNX loading error: " << e.what() << std::endl;
        throw;
    }
}

AIEnemyController::~AIEnemyController() = default;

std::vector<float> AIEnemyController::predict(
    float enemyX, float enemyY,
    float playerX, float playerY,
    float relVelX, float relVelY) {

    float inputData[6] = {
        (playerX - enemyX) / 800.0f,
        (playerY - enemyY) / 600.0f,
        enemyX / 800.0f,
        enemyY / 600.0f,
        relVelX / 100.0f,
        relVelY / 100.0f
    };

    Ort::MemoryInfo memoryInfo = Ort::MemoryInfo::CreateCpu(
        OrtArenaAllocator, OrtMemTypeDefault);
    std::vector<int64_t> inputShape = { 1, 6 };
    Ort::Value inputTensor = Ort::Value::CreateTensor<float>(
        memoryInfo, inputData, 6, inputShape.data(), inputShape.size());

    auto outputs = session->Run(
        Ort::RunOptions{ nullptr },
        inputNames.data(), &inputTensor, 1,
        outputNames.data(), 1);

    float* rawOutput = outputs.front().GetTensorMutableData<float>();
    float dx = rawOutput[0];
    float dy = rawOutput[1];

    const float smoothFactor = 0.3f;
    dx = smoothFactor * lastDx + (1 - smoothFactor) * dx;
    dy = smoothFactor * lastDy + (1 - smoothFactor) * dy;
    lastDx = dx; lastDy = dy;

    float len = std::sqrt(dx * dx + dy * dy);
    if (len > 0.001f) {
        dx /= len;
        dy /= len;
    }
    else {
        dx = dy = 0.0f;
    }

    return { dx, dy };
}